"""a module to extract and view various itk variables"""
import itk
import numpy as na

def getListFromITKObject(obj):
    try:
        data = []
        try:
            ne = obj.GetNumberOfElements()
        except:
            print "couldn't determine number of elements"
            try:
                ne = obj.GetIndexDimension()
            except:
                print "couldn't determine index dimension"
                try:
                    ne = obj.Size()
                except:
                    print "failed to determine object size"
                    try:
                        ne = obj.GetNumberOfComponents()
                    except:
                        print "failed to determine number of components"
                        return []
        for i in range(ne):
            data.append(obj.GetElement(i))
        print data
        return data
    except Exception, error:
        print "failed in getListFromITKObject", error
        return []
def print_itkPoint(obj, name = None):
    try:
        #print type(obj)
        if( name ):
            print name
        for i in range(obj.GetPointDimension()):
            print i, obj.GetElement(i)," ",
        print
    except Exception, error:
        print "failed in print_itkSize(): ", error

def print_itkIndex(obj,name=None):
    try:
        print type(obj)
        if( name ):
            print name
        for i in range(obj.GetIndexDimension()):
            print i, obj.GetElement(i)," ",
        print
    except Exception, error:
        print "failed in print_itkIndex(): ", error

def print_itkRegion(obj,name = None):
    try:
        if( name ):
            print name
        print_itkSize(obj.GetSize())
    except Exception, error:
        print "failed in print_itkRegion(): ", error
def print_itkArray(obj,name = None):
    try:
        print type(obj)
        if( name ):
            print name
        for i in range(obj.GetNumberOfElements()):
            print i, obj.GetElement(i)," ",
        print
    except Exception, error:
        print "failed in print_itkSize(): ", error
def print_itkVector(obj, name = None):
    try:
        print type(obj)
        if( name ):
            str = name[:]+": ["
        else:
            str = ""
        for i in range(obj.GetNumberOfComponents()):
            str += "%s, "%obj.GetElement(i)
        str += ']\n'
        print str
        
    except Exception, error:
        print "failed in print_itkSize(): ", error
def print_itk_variable(obj, name = None):
    try:
        print type(obj)
        if( name ):
            str = name[:]+": ["
        else:
            str = ""
        for i in range(obj.Size()):
            str += "%s, "%obj.GetElement(i)
        str += ']\n'
        print str
        
    except Exception, error:
        print "failed in print_itkSize(): ", error
def print_itkSize(obj, name = None):
    try:
        print type(obj)
        if( name ):
            str = name[:]+": ["
        else:
            str = ""
        for i in range(obj.GetSizeDimension()):
            str += "%s, "%obj.GetElement(i)
        str += ']\n'
        print str
        
    except Exception, error:
        print "failed in print_itkSize(): ", error
def get_itkVector(obj):
    try:
        v = []
        for i in range(obj.GetNumberOfComponents()):
            v.append(obj.GetElement(i))
        return v
    except Exception, error:
        print "failed in get_itkVector", error
        
def get_itkPoint(obj):
    try:
        v = []
        for i in range(obj.GetPointDimension()):
            v.append(obj.GetElement(i))
        return v
    except Exception, error:
        print "failed in get_itkVector", error

def getImgPixelDimType(img):
    return getImageType(img)

def getImageType(img):
    """determine the pixel type of the image"""
    try:
        dim = img.GetImageDimension()
        types = [itk.F,itk.SS,itk.US,itk.UC,itk.SC]
        for t in types:
            itype = itk.Image[t,dim]
            if( isinstance(img,itype)) :
                return itype
    except Exception, error:
        print "failed in getImageType", error
def numpyTypeToITKType(data):
    """compute the ITK data type corresponding to the numpy arrray data"""
    ntype = data.dtype.name
    dim = data.ndim
    if( ntype == 'uint8' ):
        return itk.Image[itk.UC,dim]
    elif( ntype == 'int8'):
        return itk.Image[itk.SC,dim]
    elif( ntype == 'uint16' ):
        return itk.Image[itk.US,dim]
    elif( ntype == 'int16' ):
        return itk.Image[itk.SS,dim]
    else:
        return itk.Image[itk.F,dim]
def getImageDim(img):
    """legacy function"""
    return img.GetImageDimension()
def getITKType(imgMode ):        
    if( imgMode.lower() == 'float' ):
        iitype = itk.F
    elif( imgMode.lower() == 'sshort' ):
        iitype = itk.SS
    elif( imgMode.lower() == 'uchar' ):
        iitype = itk.UC
    elif( imgMode.lower() == 'char' ):
        iitype = itk.C
    else:
        iitype = itk.US      
    return iitype
def getImgMode(iitype):
    if( iitype == itk.F ):
        imgMode = 'float'
    elif( iitype == itk.SS ):
        imgMode = 'sshort'
    elif( iitype == itk.UC ):
        imgMode = 'uchar'
    elif( iitype == itk.C ):
        imgMode = 'char'
    else:
        imgMode = 'ushort'
    return imgMode
    
def getImgDim(img):   
    """legacy function"""
    return img.GetImageDimension()
def getImgSize(img,returnITK=True):
    """return the size of the image. If returnITK is true return it as an itkSize object.
    Otherwise return it as a Python list"""
    try:
        lr = img.GetLargestPossibleRegion()
        sz = lr.GetSize()
        if( returnITK ):
            return sz
        dim = []
        for i in range(sz.GetSizeDimension()):
            dim.append(sz.GetElement(i))
        return dim
    except Exception, error:
        print "failed in getImgSize()", error
def get_itkSize(obj):
    try:
        sz = []
        for i in range(obj.GetSizeDimension()):
            sz.append(obj.GetElement(i))
        return sz
    except Exception, error:
        print "failed in get_itkSize()", error
def getITKDirection(obj):
    try:
        d = obj.GetDirection()
        dv = d.GetVnlMatrix()
        dir = na.zeros((dv.rows(), dv.columns()), na.float32)
        for i in range(dv.rows()):
            for j in range(dv.columns()):
                dir[i, j] = dv.get(i, j)
        return dir
    except Exception,  error:
        print "failed in getITKDirections",  error
def setHeaderData(img,refimg):
    """set the header information in img based on that contained in refimg"""
    try:

       
        img.SetOrigin(refimg.GetOrigin())

        img.SetSpacing(refimg.GetSpacing())

        img.SetDirection(refimg.GetDirection())
        return img
    except Exception, error:
        print "failed in setHeaderData", error
        return img             
        
def getObjectSize(img):
    try:
        try:
            dim = img.GetSpacing().GetNumberOfComponents()
            
        except Exception, error:
            print "failed in getCaster", error
            try:
                dim = img.GetSize()
            except Exception, error:
                print "failed in getCaster", error
                try:
                    dim = img.GetImageDimension()
                except Exception, error:
                    print "failed in getCaster", error
                    dim = input("Enter image dimension\n")  
        return dim
    except Exception, error:
        print "failed in getObjectSize()", error
        return None    
        
        
        
def printImageData(img,label=None):
    try:
        print "*"*42
        if( label ):
            print label
        print_itk_variable(img.GetOrigin(),name="Image Origin") 
        print_itkVector(img.GetSpacing(),name="Image Spacing") 
        d = img.GetDirection()
        d2 = d.GetVnlMatrix()
        dir = 'dir='
        try:
                        for i in range(d2.rows()):
                                for j in range(d2.columns()):
                                        dir +='%s, '%d2.get(i,j)
        except:
                pass
        print dir
        print "Image Range: ",itk.range(img)            
        region = img.GetLargestPossibleRegion()
        sz = region.GetSize()  
        s = 'size='
        for i in range(region.GetImageDimension()):  
            s += '%3d,'%sz.GetElement(i)
        print s
        print "\n"
    except Exception, error:
        print "failed in printImageData", error
                
def getPixelValue(img,crds, iType=itk.Image[itk.F,3] ):
    """gets the Pixel value at x,y,z. Currently have to convert to array"""
    try:
        array = itk.PyBuffer[iType].GetArrayFromImage(img)
        return array[crds[2],crds[1],crds[0]]
    except Exception, error:
        print "faile din getPixelValue", error
